#include "vdsse.h"

// 初始化
void VDSSE_Setup(int lambda, nlohmann::json& clientJson) {
    assert(!(lambda & 3));
    using namespace nlohmann;
    gmp_randstate_t grt;
    gmp_randinit_default(grt); //设置随机数生成算法为默认
    gmp_randseed_ui(grt, time(NULL));	//设置随机化种子为当前时间
    json &K = clientJson["K"];
    for (json::iterator it = K.begin(); it != K.end(); ++it) {
        mpz_t key;
        mpz_init(key);
        mpz_urandomb(key, grt, lambda);
        it.value() = mpz_class(key).get_str(16);
        mpz_clear(key);
    }
}

void Initialization(nlohmann::json& clientJson) {
    using namespace nlohmann;
    clientJson["V"] = 1;
    json& files = clientJson["files"];
    json& W = clientJson["W"];
    json& K = clientJson["K"];
    std::unordered_set<std::string> wset;
    // 初始化后面要用到的几个 PRF，注意只通过单例模板调用
    int l = 128, n = 256, logN = 128;
    mpz_class k_prp(clientJson["K"]["Kprp"], 16);
    mpz_class k_prf(clientJson["K"]["Kprf"], 16);
    mpz_class k_p(clientJson["K"]["Kp"], 16);
    PRF &pie = Singleton<PiePrf>::Instance();
    pie.setup(k_prp, l);
    PRF &f = Singleton<FPrf>::Instance();
    f.setup(k_prf, n);
    PRF &P = Singleton<PPrf>::Instance();
    P.setup(k_p, logN);

    for (auto it = files.begin(); it != files.end(); it++) {
        (*it)["v"] = 1;
        json &M = (*it)["M"];
        M.clear();
        std::string Cj = SKE_Enc(mpz_class(K["Kenc"], 16), (*it)["text"]);
        (*it)["C"] = Cj;
        int Cj_len = Cj.size();
        int i = 0;
        while (i + 256 < Cj_len) {
            M.push_back(Cj.substr(i, i + 256));
            i += 256;
        }
        M.push_back(Cj.substr(i));

        (*it)["alpha"] = P((*it)["text"]);
        json &keywords = (*it)["keywords"];
        for (auto key_it = keywords.begin(); key_it != keywords.end(); key_it++) {
            wset.insert(*key_it);
        }
    }
    W.clear();
    for (auto& ele: wset) {
        W.push_back(ele);
    }
}

void BuildST(nlohmann::json& clientJson, ST& searchTable) {
    using namespace nlohmann;
    json& files = clientJson["files"];
    json& W = clientJson["W"];
    json& ST = clientJson["ST"];
    ST.clear();
    PRF pie = Singleton<PiePrf>::Instance();
    SHA256 hashfunc;
    for (auto& w: W) {
        std::string wi = w.get<std::string>();
        ST.push_back(pie(wi));
        mpz_class Kw(hashfunc(wi), 16);
        for (auto& file: files) {
            std::string wij = "0";
            if (file["keywords"].contains(w)) {
                wij = "1";
            }
            ST.push_back(SKE_Enc(Kw, wij + std::to_string(int(file["v"]))));
        }
    }
    int n = W.size(), N = files.size() + 1;
    searchTable.setup(n, N);
    int ind = 0;

    Node* ptr = nullptr;
    Node* upptr = nullptr;
    for (auto it = ST.begin(); it != ST.end(); it++, ind++) {
        if ((ind % N) == 0) {
            searchTable.Lw[ind / 4] = new Node((*it).get<std::string>());
            ptr = searchTable.Lw[ind / 4];
            if (ind >= N) {
                upptr = searchTable.Lw[ind / 4 - 1];
                upptr->down = ptr;
            }
        } else {
            ptr->right = new Node((*it).get<std::string>());
            ptr = ptr->right;
            if (upptr) {
                upptr = upptr->right;
                upptr->down = ptr;
            }
        }
    }
    ptr = searchTable.Lw[0];
    for (int it = 0; it < N; it++) {
        searchTable.Lf[it] = ptr;
        ptr = ptr->right;
    }
}

// 要先使 verificationList 指向一个空的 Node 再传参
void BuildVL(nlohmann::json& clientJson, Node* verificationList) {
    using namespace nlohmann;
    using std::string;
    json& files = clientJson["files"];
    json& W = clientJson["W"];
    Node* ptr = verificationList;
    PRF f = Singleton<FPrf>::Instance();
    PRF pie = Singleton<PiePrf>::Instance();
    for (auto& w: W) {
        mpz_class f_v(f(clientJson["V"].get<string>()), 16);
        mpz_class f_pie_w(f(pie(w.get<string>())), 16);
        mpz_class AAT = f_v + f_pie_w;
        for (auto& file: files) {
            if (file["keywords"].contains(w)) {
                mpz_class alpha(file["alpha"].get<string>(), 16);
                for (auto& m: file["M"]) {
                    AAT += alpha * mpz_class(m.get<string>(), 16);
                }
            }
        }
        ptr->right = new Node(AAT.get_str(16));
        ptr = ptr->right;
    }
}

void IndexBuild(nlohmann::json& clientJson, ST& searchTable, Node* verificationList) {
    Initialization(clientJson);
    BuildST(clientJson, searchTable);
    BuildVL(clientJson, verificationList);
}

Token GenToken(std::string w) {
    using std::string;
    PRF pie = Singleton<PiePrf>::Instance();
    SHA256 hashfunc;
    return Token(pie(w), hashfunc(w));
}

void Search(nlohmann::json& clientJson, Token Tw, ST& searchTable, Node* verificationList, std::vector<int>& Cw_id, std::string& ATT) {
    using std::string;
    nlohmann::json& files = clientJson["files"];
    Cw_id.clear();
    string n1 = Tw.pie_w;
    mpz_class theta(Tw.Kw, 16);
    Node* lwk = nullptr;
    int k;
    for (int i = 0; i < searchTable.n_; i++) {
        lwk = searchTable.Lw[i];
        if (strcmp(lwk->data.c_str(), n1.c_str()) == 0) {
            k = i + 1;
            break;
        }
    }
    if (lwk == nullptr) return;
    Node* ptr = lwk->right;
    for (int i = 1; i < searchTable.N_; i++, ptr = ptr->right) {
        string wv = SKE_Dec(theta, ptr->data);
        assert(strcmp(wv.c_str(), "0") != 0);
        if (wv[0] == '1') {
            Cw_id.push_back(i - 1);
        }
    }
    for (int i = 0; i < k; i++) {
        verificationList = verificationList->right;
    }
    ATT = verificationList->data;
}

bool Verify(nlohmann::json& clientJson, Token Tw, std::vector<int>& Cw_id, std::string AAT) {
    using std::string;
    nlohmann::json& files = clientJson["files"];
    PRF f = Singleton<FPrf>::Instance();
    PRF P = Singleton<PPrf>::Instance();
    mpz_class f_n1(f(Tw.pie_w), 16);
    mpz_class f_V(f(clientJson["V"].get<string>()), 16);
    mpz_class AAT_1 = f_n1 + f_V;
    for (auto it = Cw_id.begin(); it != Cw_id.end(); it++) {
        int j = *it;
        nlohmann::json& file = files[j];
        mpz_class alpha(P(file["text"]), 16);
        for (auto& m: file["M"]) {
            AAT_1 += alpha * mpz_class(m.get<string>(), 16);
        }
    }
    return AAT_1 == mpz_class(AAT, 16);
}

